# 导入包
import pandas as pd
import statsmodels.formula.api as smf
from sklearn.model_selection import train_test_split
from sqlalchemy import create_engine
from sklearn.metrics import r2_score, mean_squared_error
from matplotlib import pyplot as plt
import statsmodels.api as sm
import numpy as np
import seaborn as sns
import pymysql

db_config = {
    'host': 'localhost',
    'user': 'root',
    'password': '123456',
    'database': 'ljy',
    'port': 3306,          # 明确指定端口
    'charset': 'utf8mb4'   # 添加字符集设置
}
engine = create_engine(
    f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}"
)
conn = pymysql.connect(**db_config)
chunk_size = 10000

# 获取中国平安日线数据
df = pd.read_sql_query(
        """
        SELECT d.*, i.closes as i_closes, i.vol as i_vol
        FROM date_1 d
        WHERE d.trade_date BETWEEN '2023-01-01' AND '2023-12-31' AND d.ts_code='000001.SZ'
        """, 
        conn, 
        chunksize=chunk_size
    )
df1 = pd.concat(df, ignore_index=True)

df1.head()
# 增加一列，股票的涨跌幅
df1['zd_close'] = round((df1['closes']-df1['closes'].shift(1))/df1['closes'].shift(1),2)
df1['hz_close'] = round((df1['i_closes']-df1['i_closes'].shift(1))/df1['i_closes'].shift(1),2)
df1['hz_vol'] = round((df1['i_vol'].shift(1)-df1['i_vol'].shift(2))/df1['i_vol'].shift(2),2)
# df1['vol1'] = round((df1['i_vol'])/df1['vol'],6)

# 处理缺失值
df1 = df1.dropna(subset=['zd_close'])

# 划分训练集与测试集（随机划分）
train_df, test_df = train_test_split(df1, test_size=0.2, random_state=42)

# 在训练集上筛选特征
numeric_cols_train = train_df.select_dtypes(include=['number']).columns.tolist()
# numeric_cols = df1.select_dtypes(include=['number']).columns.tolist()
# 筛选合适的自变量
excluded = ['zd_close', 'ts_code', 'trade_date', 'id', 'pre_closes', 'changes', 'pct_chg', 'opens', 'high', 'low', 'closes', 'i_closes', 'amount', 'i_vol']  # 根据实际列名调整
predictors_train  = [col for col in numeric_cols_train if col not in excluded]

# 排除方差为零的列 移除所有值相同的列，避免共线性问题
predictors_train  = [col for col in predictors_train  if df1[col].nunique() > 1]

# 检查样本量与自变量数量
n_train = len(train_df)
k_train = len(predictors_train)
if n_train <= k_train + 1:
    print(f"错误：训练样本数{n_train}不足，至少需要{k_train + 2}个样本。")

else:
     # 训练模型
    formula = 'zd_close ~ ' + ' + '.join(predictors_train)
    results1 = smf.ols(formula, data=train_df).fit()
    print("==== 训练集模型结果 ====")
    print(results1.summary())
    
    # 测试集评估
    X_test = test_df[predictors_train]
    y_test = test_df['zd_close']
    y_pred = results1.predict(X_test)
    
    # 计算评估指标
    print("\n==== 测试集评估结果 ====")
    print(f"R²分数: {r2_score(y_test, y_pred):.4f}")
    print(f"均方误差: {mean_squared_error(y_test, y_pred):.4f}")
    